from torch import nn


class Model(nn.Module):
    in_channels:int
    data_size:int


    def __init__(self, dst:str):
        super(Model,self).__init__()
        self._match_dst(dst)
        self._get_fc1_in_features_size()


    def _match_dst(self,dst:str):
        if dst == "mnist":
            self.in_channels = 1
            self.data_size = 28
        elif dst == "cifar10":
            self.in_channels = 3
            self.data_size = 32
        else:
            assert False,"数据集名称错误"


    def _get_fc1_in_features_size(self):
        #输出尺寸 = (输入尺寸 - 卷积核尺寸 + 2 * 填充) / 步长 + 1
        #out_size = (in_size - kernel_size + 2 * padding) / stride + 1
        raise NotImplemented
    